package server

import (
	"bufio"
	"encoding/binary"
	"fmt"
	"go-socks5/pkg/config"
	"go-socks5/pkg/log"
	"io"
	"net"
	"time"
)

const (
	SocksVer = 0x05
	AuthVer  = 0x01
)

const (
	MethodNoAuth       = 0x00
	MethodGSSAPI       = 0x01
	MethodUserPwd      = 0x02
	MethodNoAcceptable = 0xFF
)

const (
	CMDConnect      = 0x01
	CMDBind         = 0x02
	CMDUDPAssociate = 0x03
)

const (
	ATYPIPv4   = 0x01
	ATYPDomain = 0x03
	ATYPIPv6   = 0x04
)

const (
	ReplySuccess                 = 0x00
	ReplyServerFailure           = 0x01
	ReplyConnectionNotAllowed    = 0x02
	ReplyNetworkUnreachable      = 0x03
	ReplyHostUnreachable         = 0x04
	ReplyConnectionRefused       = 0x05
	ReplyTTLExpired              = 0x06
	ReplyCommandNotSupported     = 0x07
	ReplyAddressTypeNotSupported = 0x08
	ReplyUnassigned              = 0x09
)

type Socks5Server struct {
	Config        *config.ServerConfig
	Listener      net.Listener
	ClientManager *ClientManager
}

type Socks5Client struct {
	Conn       net.Conn // 客户端链接
	TargetConn net.Conn // 目标服务器链接
	State      int      // 客户端当前状态 0.握手阶段 1.认证阶段 2.请求阶段 3.relay阶段
	closeChan  chan error
	Username   string
	Password   string
	reader     *bufio.Reader
}

func NewSocks5Client(conn net.Conn, username string, password string, reader *bufio.Reader) *Socks5Client {
	return &Socks5Client{
		Conn:      conn,
		closeChan: make(chan error, 1),
		State:     0,
		Username:  username,
		Password:  password,
		reader:    reader,
	}
}

func NewSocks5Server(config *config.ServerConfig) (*Socks5Server, error) {
	addr := fmt.Sprintf("%s:%d", config.BindAddr, config.BindPort)
	listener, err := net.Listen("tcp", addr)
	if err != nil {
		return nil, fmt.Errorf("failed to listen on %s: %v", addr, err)
	}

	return &Socks5Server{
		Config:        config,
		Listener:      listener,
		ClientManager: NewClientManager(),
	}, nil
}

func (s *Socks5Server) Run() error {
	for {
		conn, err := s.Listener.Accept()
		if err != nil {
			return fmt.Errorf("failed to accept connection: %v", err)
		}
		log.Infof("new client connection: %v", conn.RemoteAddr())
		go s.handleConnection(conn)
	}
}

func (s *Socks5Server) handleConnection(conn net.Conn) {
	clientAddr := conn.RemoteAddr().String()

	defer func() {
		log.Infof("[%s] client disconnected", clientAddr)
		s.ClientManager.RemoveClient(clientAddr)
		conn.Close()
	}()

	reader := bufio.NewReader(conn)

	// 创建客户端实例
	client := NewSocks5Client(conn, s.Config.Username, s.Config.Password, reader)

	// 保存客户端信息
	s.ClientManager.AddClient(&ClientInfo{
		ClientAddr:  clientAddr,
		ConnectTime: time.Now(),
	})

	// 初始阶段：协商
	err := handleConsultWith(client)
	if err != nil {
		log.Errorf("[%s] consultation failed: %v", clientAddr, err)
		return
	}

	// 根据协商结果决定下一步
	log.Infof("[%s] client enter auth state: %d", clientAddr, client.State)
	if client.State == 1 {
		// 需要认证
		err = handleAuth(client)
		if err != nil {
			log.Errorf("[%s] authentication failed: %v", clientAddr, err)
			return
		}
	}

	// 处理请求
	log.Infof("[%s] client enter request state: %d", clientAddr, client.State)
	if client.State == 2 {
		err = handleRequest(client)
		if err != nil {
			log.Errorf("[%s] request failed: %v", clientAddr, err)
			return
		}
	}

	// 数据转发阶段
	log.Infof("[%s] client enter relay state: %d", clientAddr, client.State)
	if client.State == 3 && client.TargetConn != nil {
		// 更新客户端信息
		s.ClientManager.UpdateClientStats(clientAddr, client.TargetConn.RemoteAddr().String(), client.State)

		log.Infof("[%s] entering relay mode for %s <-> %v",
			clientAddr,
			clientAddr,
			client.TargetConn.RemoteAddr())

		// 一旦进入数据转发模式，就持续传输数据直到连接关闭
		// 这是 SOCKS5 协议的终止状态，不会有更多的命令
		err := handleRelay(client, s.ClientManager)
		if err != nil {
			log.Errorf("[%s] relay failed: %v", clientAddr, err)
			return
		}
	}

	// 如果没有达到预期的状态，结束连接处理
	if client.State != 3 {
		log.Errorf("[%s] unexpected client state: %d", clientAddr, client.State)
		return
	}
}

// handleConsultWith 协商阶段包处理
// request:
// +----+----------+----------+
// |VER | NMETHODS | METHODS  |
// +----+----------+----------+
// | 1  |    1     | 1 to 255 |
// +----+----------+----------+
//
// response:
// +----+--------+
// |VER | METHOD |
// +----+--------+
// | 1  |   1    |
// +----+--------+
func handleConsultWith(client *Socks5Client) error {
	reader := client.reader

	// 读取版本号
	version, err := reader.ReadByte()
	if err != nil || version != SocksVer {
		return fmt.Errorf("failed to read version: %v", err)
	}

	// 读取认证方法数量
	nmethods, err := reader.ReadByte()
	if err != nil {
		return fmt.Errorf("failed to read number of methods: %v", err)
	}

	// 读取认证方法列表
	methods := make([]byte, nmethods)
	_, err = reader.Read(methods)
	if err != nil {
		return fmt.Errorf("failed to read methods: %v", err)
	}

	// 选择认证方法
	// 发送认证方法选择响应
	if client.Username != "" && client.Password != "" {
		_, err = client.Conn.Write([]byte{SocksVer, MethodUserPwd})
		if err != nil {
			return fmt.Errorf("failed to write auth response: %v", err)
		}
		// 客户端通过协商阶段，进入认证阶段
		client.State = 1
	} else {
		_, err = client.Conn.Write([]byte{SocksVer, MethodNoAuth})
		if err != nil {
			return fmt.Errorf("failed to write auth response: %v", err)
		}
		// 客户端通过协商阶段，进入请求阶段
		client.State = 2
	}

	return nil
}

// handleAuth 处理认证阶段数据包
// +----+------+----------+------+----------+
// |VER | ULEN |  UNAME   | PLEN |  PASSWD  |
// +----+------+----------+------+----------+
// | 1  |  1   | 1 to 255 |  1   | 1 to 255 |
// +----+------+----------+------+----------+
//
// +----+--------+
// |VER | STATUS |
// +----+--------+
// | 1  |   1    |
// +----+--------+
func handleAuth(client *Socks5Client) error {
	reader := client.reader

	// 读取版本号
	version, err := reader.ReadByte()
	if err != nil || version != AuthVer {
		return fmt.Errorf("failed to read version: %v", err)
	}

	// 读取用户名长度
	ulen, err := reader.ReadByte()
	if err != nil {
		return fmt.Errorf("failed to read username length: %v", err)
	}

	// 读取用户名
	username := make([]byte, ulen)
	_, err = reader.Read(username)
	if err != nil {
		return fmt.Errorf("failed to read username: %v", err)
	}

	// 读取密码长度
	plen, err := reader.ReadByte()
	if err != nil {
		return fmt.Errorf("failed to read password length: %v", err)
	}

	// 读取密码
	password := make([]byte, plen)
	_, err = reader.Read(password)
	if err != nil {
		return fmt.Errorf("failed to read password: %v", err)
	}

	// 获取配置中的用户名和密码
	if client.Username == string(username) && client.Password == string(password) {
		// 发送认证成功响应
		_, err = client.Conn.Write([]byte{AuthVer, ReplySuccess})
		if err != nil {
			return fmt.Errorf("failed to write auth response: %v", err)
		}
		// 客户端通过认证阶段，进入请求阶段
		client.State = 2
	} else {
		client.Conn.Write([]byte{AuthVer, ReplyServerFailure})
		return fmt.Errorf("auth failed, wrong username or password")
	}

	return nil
}

// handleRequest 处理客户端请求阶段数据包
// socks5 client request
// +----+-----+-------+------+----------+----------+
// |VER | CMD |  RSV  | ATYP | DST.ADDR | DST.PORT |
// +----+-----+-------+------+----------+----------+
// | 1  |  1  | X'00' |  1   | Variable |    2     |
// +----+-----+-------+------+----------+----------+
//
// socks5 server response
// +----+-----+-------+------+----------+----------+
// |VER | REP |  RSV  | ATYP | BND.ADDR | BND.PORT |
// +----+-----+-------+------+----------+----------+
// | 1  |  1  | X'00' |  1   | Variable |    2     |
// +----+-----+-------+------+----------+----------+
func handleRequest(client *Socks5Client) error {
	reader := client.reader

	// 读取版本号
	version, err := reader.ReadByte()
	if err != nil || version != SocksVer {
		return fmt.Errorf("failed to read version: %v", err)
	}

	cmd, err := reader.ReadByte()
	if err != nil {
		return fmt.Errorf("failed to read cmd: %v", err)
	}

	if cmd != CMDConnect && cmd != CMDBind && cmd != CMDUDPAssociate {
		return fmt.Errorf("illegal cmd: %d", cmd)
	}

	// rsv 保留字段
	_, err = reader.ReadByte()
	if err != nil {
		return fmt.Errorf("failed to read rsv: %v", err)
	}

	// 读取目标地址类型
	atyp, err := reader.ReadByte()
	if err != nil {
		return fmt.Errorf("failed to read atyp: %v", err)
	}

	var dstAddr string
	switch atyp {
	case ATYPIPv4:
		// 读取4字节的IPv4地址
		ipv4 := make([]byte, 4)
		if _, err = io.ReadFull(reader, ipv4); err != nil {
			return fmt.Errorf("failed to read IPv4 address: %v", err)
		}
		// 转换为点分十进制格式
		dstAddr = net.IPv4(ipv4[0], ipv4[1], ipv4[2], ipv4[3]).String()
		// 稍后使用 dstAddr 连接目标服务器
	case ATYPDomain:
		// 读取域名长度
		domainLen, err := reader.ReadByte()
		if err != nil {
			return fmt.Errorf("failed to read domain length: %v", err)
		}
		// 读取域名
		domain := make([]byte, domainLen)
		if _, err = io.ReadFull(reader, domain); err != nil {
			return fmt.Errorf("failed to read domain: %v", err)
		}
		dstAddr = string(domain)
	case ATYPIPv6:
		// 读取16字节的IPv6地址
		ipv6 := make([]byte, 16)
		if _, err = io.ReadFull(reader, ipv6); err != nil {
			return fmt.Errorf("failed to read IPv6 address: %v", err)
		}
		// 转换为冒号分隔的格式
		dstAddr = net.IP(ipv6).String()
	}

	// 读取目标地址
	port := make([]byte, 2)
	if _, err = io.ReadFull(reader, port); err != nil {
		return fmt.Errorf("failed to read port: %v", err)
	}

	// 将两个字节转换为uint16端口号
	dstPort := binary.BigEndian.Uint16(port)

	// 使用 dstAddr 和 dstPort 连接目标服务器
	var targetAddr string
	if atyp == byte(ATYPIPv6) {
		targetAddr = fmt.Sprintf("[%s]:%d", dstAddr, dstPort)
	} else {
		targetAddr = fmt.Sprintf("%s:%d", dstAddr, dstPort)
	}

	log.Infof("[%s] cmd: %d, targetAddr: %v", client.Conn.RemoteAddr(), cmd, targetAddr)

	var targetConn net.Conn
	switch cmd {
	case CMDConnect:
		// targetConn, err = net.DialTimeout("tcp", targetAddr, time.Second*3)
		targetConn, err = net.Dial("tcp", targetAddr)
		if err != nil {
			return fmt.Errorf("failed to use tcp connect to target server: %v", err)
		}
	case CMDUDPAssociate:
		targetConn, err = net.Dial("udp", targetAddr)
		if err != nil {
			return fmt.Errorf("failed to use udp connect to target server: %v", err)
		}
	}

	log.Infof("[%s] connected to targetAddr: %v", client.Conn.RemoteAddr(), targetAddr)

	// 发送响应
	// 当 relay server 和 socks5 server 是同一台服务器时，BND.ADDR 和 BND.PORT 的值全部为 0 即可
	response := []byte{SocksVer, ReplySuccess, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00}
	_, err = client.Conn.Write(response)
	if err != nil {
		return fmt.Errorf("failed to write response: %v", err)
	}

	// 保存目标服务器链接
	client.TargetConn = targetConn

	// 客户端通过请求阶段，进入 relay 阶段
	client.State = 3

	return nil
}

// CountReader 是一个可以计数的 Reader 包装器
type CountReader struct {
	io.Reader
	BytesRead int64
	OnRead    func(n int64) // 回调函数，在每次读取后调用
}

// Read 重写 Read 方法，计数并调用回调
func (cr *CountReader) Read(p []byte) (n int, err error) {
	n, err = cr.Reader.Read(p)
	cr.BytesRead += int64(n)
	if cr.OnRead != nil {
		cr.OnRead(int64(n))
	}
	return
}

// handleRelay 处理客户端和目标服务器之间的数据转发
func handleRelay(client *Socks5Client, clientManager *ClientManager) error {
	// start proxying
	errCh := make(chan error, 2)

	clientReader := &CountReader{
		Reader: client.reader,
		OnRead: func(n int64) {
			// 更新上传流量统计
			clientManager.UpdateClientBytesSent(client.Conn.RemoteAddr().String(), n)
		},
	}

	targetReader := &CountReader{
		Reader: client.TargetConn,
		OnRead: func(n int64) {
			// 更新下载流量统计
			clientManager.UpdateClientBytesReceived(client.Conn.RemoteAddr().String(), n)
		},
	}

	go proxy(client.TargetConn, clientReader, errCh)
	go proxy(client.Conn, targetReader, errCh)

	// Wait
	for i := 0; i < 2; i++ {
		e := <-errCh
		if e != nil {
			// return from this function closes target (and conn).
			return e
		}
	}
	return nil
}

// proxy is used to suffle data from src to destination, and sends errors
// down a dedicated channel
func proxy(dst io.Writer, src io.Reader, errCh chan error) {
	_, err := io.Copy(dst, src)
	if tcpConn, ok := dst.(*net.TCPConn); ok {
		tcpConn.CloseWrite()
	}
	errCh <- err
}
